import numpy as np
import os
import cv2

file_path = "/disk527/sdb1/a804_cbf/datasets/collect_data/2024-09-14_16-38-56/events.csv"
t_window = 1e-2
height = 512
width = 512


def event_accquire(str):
    data = str.split(",")
    return int(data[0]), int(data[1]), float(data[2]), int(data[3])


def read_csv(file_path):
    with open(file_path, "r") as f:
        # Skip the first line
        data = f.readline()
        t0 = 0
        events = {"x": [], "y": [], "t": [], "p": []}
        length = 0
        for data in f:
            x, y, t, p = event_accquire(data)
            if t < 50:
                continue
            if t - t0 <= t_window:
                events["x"].append(x)
                events["y"].append(y)
                events["t"].append(t)
                events["p"].append(p)
                length += 1
            else:
                if length > 0:
                    yield events
                events = {"x": [x], "y": [y], "t": [t], "p": [p]}
                t0 = t
                length = 0


def visualizer(events):
    img = np.zeros((height, width, 3), dtype=np.uint8)
    x = np.array(events["x"])
    y = np.array(events["y"])
    p = np.array(events["p"])
    img[y, x, 1] = p * 255
    img[y, x, 2] = (1 - p) * 255
    return img


if __name__ == "__main__":
    if not os.path.exists("images"):
        os.mkdir("images")
    for i, events in enumerate(read_csv(file_path)):
        img = visualizer(events)
        t = events["t"][0]
        cv2.imwrite(f"images/{t:.3f}.png", img)
        print(f"Save image {i}.png")
